import asyncio
import time
from typing import Awaitable

from langchain.callbacks import AsyncIteratorCallbackHandler
from langchain.chains import LLMChain
from langchain_core.prompts import HumanMessagePromptTemplate, ChatPromptTemplate

from config.model_config import get_chat_openai_xin


async def chat_iterator():
    model = get_chat_openai_xin(streaming=True)

    human_template = "{text}"
    human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
    print(human_message_prompt)

    chat_prompt = ChatPromptTemplate.from_messages([human_message_prompt])
    start_time = time.time()
    query = '写一篇人工智能的文字，不少于200字'

    chain = LLMChain(
        llm=model,
        prompt=chat_prompt,
        # output_parser=CommaSeparatedListOutputParser()
    )
    callback = AsyncIteratorCallbackHandler()
    task = asyncio.create_task(
        chain.acall({"text": query}, callbacks=[callback])
    )

    print('xxxx')
    async for token in callback.aiter():
        print(token, end='', flush=True)
    await task


if __name__ == '__main__':
    asyncio.run(chat_iterator())

